{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "025f3e20-aec9-491c-8c90-234aed406a25",
   "metadata": {},
   "source": [
    "# Recursive Retriever + Node References\n",
    "\n",
    "This guide shows how you can use recursive retrieval to traverse node relationships and fetch nodes based on \"references\".\n",
    "\n",
    "Node references are a powerful concept. When you first perform retrieval, you may want to retrieve the reference as opposed to the raw text. You can have multiple references point to the same node.\n",
    "\n",
    "In this guide we explore some different usages of node references:\n",
    "- **Chunk references**: Different chunk sizes referring to a bigger chunk\n",
    "- **Metadata references**: Summaries + Generated Questions referring to a bigger chunk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "273f38de-e79a-4ce2-ad4e-2c70afc33f34",
   "metadata": {},
   "source": [
    "## Load Data + Setup\n",
    "\n",
    "In this section we download the Llama 2 paper and create an initial set of nodes (chunk size 1024)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb829ef-b54b-4095-a832-6d1d115aa645",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget --user-agent \"Mozilla\" \"https://arxiv.org/pdf/2307.09288.pdf\" -O \"data/llama2.pdf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6cd97455-5ff3-43ee-8222-f496ec234dc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from llama_hub.file.pdf.base import PDFReader\n",
    "from llama_index.response.notebook_utils import display_source_node\n",
    "from llama_index.retrievers import RecursiveRetriever\n",
    "from llama_index.query_engine import RetrieverQueryEngine\n",
    "from llama_index import VectorStoreIndex, ServiceContext\n",
    "from llama_index.llms import OpenAI\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a07c0e42-1ae8-4267-9355-6bb75323f82a",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = PDFReader()\n",
    "docs0 = loader.load_data(file=Path(\"./data/llama2.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "493e5492-a6ae-4e3e-aa23-274c0605b165",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import Document\n",
    "\n",
    "doc_text = \"\\n\\n\".join([d.get_content() for d in docs0])\n",
    "docs = [Document(text=doc_text)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7c2abcd3-6cae-49dd-8719-9b738d000652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.node_parser import SimpleNodeParser\n",
    "from llama_index.schema import IndexNode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "91b997ae-9260-4ae7-af2f-0f8d38625d32",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_parser = SimpleNodeParser.from_defaults(chunk_size=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0cda44b0-fd27-4255-9aa7-08d358635772",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_nodes = node_parser.get_nodes_from_documents(docs)\n",
    "# set node ids to be a constant\n",
    "for idx, node in enumerate(base_nodes):\n",
    "    node.id_ = f\"node-{idx}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aeacfed8-ccaf-41fb-aed6-5cca437c8728",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f43ebab2-fc46-41ea-8a92-9148994d793f",
   "metadata": {},
   "source": [
    "## Baseline Retriever\n",
    "\n",
    "Define a baseline retriever that simply fetches the top-k raw text nodes by embedding similarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "704fb3da-710e-4ad9-b630-565911917f0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_index = VectorStoreIndex(base_nodes)\n",
    "base_retriever = base_index.as_retriever(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "160c339b-601a-486b-9e17-dd6cc9f133ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrievals = base_retriever.retrieve(\n",
    "    \"Can you tell me about the key concepts for safety finetuning\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "632610f3-c8f2-440a-ab27-5ca7d65f882a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Node ID:** node-22<br>**Similarity:** 0.8322523818232979<br>**Text:** Wethendefinebestpracticesforsafeandhelpfulmodelresponses: themodelshouldfirstaddressimmediate\n",
       "safetyconcernsifapplicable,thenaddressthepromptbyexplainingthepotentialriskstotheuser,andfinally\n",
       "provide additional information if possible.We also ask the annotators to avoid negative user experience\n",
       "categories (see Appendix A.5.2).The guidelines are meant to be a general guide for the model and are\n",
       "iteratively refined and revised to include newly identified risks.4.2.2 Safety Supervised Fine-Tuning\n",
       "InaccordancewiththeestablishedguidelinesfromSection4.2.1,wegatherpromptsanddemonstrations\n",
       "ofsafemodelresponsesfromtrainedannotators,andusethedataforsupervisedfine-tuninginthesame\n",
       "manner as described in Section 3.1.An example can be found in Table 5.The annotators are instructed to initially come up with prompts that they think could potentially induce\n",
       "themodel toexhibit unsafebehavior, i.e.,perform redteaming, asdefined bythe guidelines.Subsequently,\n",
       "annotators are tasked with crafting a safe and helpful response that the model should produce.4.2.3 Safety RLHF\n",
       "Weobserveearlyinthedevelopmentof Llama 2-Chat thatitisabletogeneralizefromthesafedemonstrations\n",
       "insupervisedfine-tuning.Themodelquicklylearnstowritedetailedsaferesponses,addresssafetyconcerns,\n",
       "explainwhythetopicmightbesensitive,andprovideadditionalhelpfulinformation.Inparticular,when\n",
       "the model outputs safe responses, they are often more detailed than what the average annotator writes.Therefore, after gathering only a few thousan...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Node ID:** node-18<br>**Similarity:** 0.79782506307983<br>**Text:** Next,wedescribe\n",
       "theprocessofoursafetyalignment(Section4.2),explaininghowwecollectedsafety-relatedannotationsand\n",
       "utilizedSFTandRLHF,andpresentexperimentalresults.Then,wediscusstheredteamingweperformedto\n",
       "furtherunderstandandimprovemodelsafety(Section4.3).Finally,wepresentquantitativesafetyevaluations\n",
       "ofLlama 2-Chat (Section 4.4).We also share a model card in the Appendix, in Table 52.4.1 Safety in Pretraining\n",
       "It is important to understand what is in the pretraining data both to increase transparency and to shed\n",
       "lightonrootcausesofpotentialdownstreamissues,suchaspotentialbiases.Thiscaninformwhat,ifany,\n",
       "downstream mitigations to consider, and help guide appropriate model use.In this section, we analyze the\n",
       "pretraining datafor distributionsof languages,demographic representations,and toxicity.Wealso present\n",
       "the results of testing the pretrained models on existing safety benchmarks.Steps Taken to Pretrain Responsibly.We followed Meta’s standard privacy and legal review processes for\n",
       "each dataset used in training.We did not use any Meta user data in training.We excluded data from certain\n",
       "sitesknowntocontainahighvolumeofpersonalinformationaboutprivateindividuals.Wemadeabest\n",
       "effort to train our models efficiently to reduce the carbon footprint of pretraining (Section 2.2.1).Sharing our\n",
       "modelsbroadlywillreducetheneedforotherstotrainsimilarmodels.Noadditionalfilteringwasconducted\n",
       "onthedatasets,toallow Llama 2 tobemorewidelyusableacrosstasks(e.g.,itcanbebetterusedforhate\n",
       "speechclassif...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for n in retrievals:\n",
    "    display_source_node(n, source_length=1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "96dd8a01-1cae-4614-beab-5b5e0434fefe",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine_base = RetrieverQueryEngine.from_args(\n",
    "    base_retriever, service_context=service_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "82ae66ff-7d12-45c8-9b1a-adb20bd3c7ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The key concepts for safety fine-tuning include addressing immediate safety concerns, explaining potential risks to the user, and providing additional information if possible. The guidelines for safety fine-tuning are meant to be a general guide for the model and are iteratively refined and revised to include newly identified risks. Safety fine-tuning can be done through supervised fine-tuning, where prompts and demonstrations of safe model responses are gathered from trained annotators, or through reinforcement learning from human feedback (RLHF), where human preference data is collected to train a safety reward model. The goal of safety fine-tuning is to make the model more robust to unsafe behavior and improve its ability to generate nuanced and detailed safe responses.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine_base.query(\n",
    "    \"Can you tell me about the key concepts for safety finetuning\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5431df3-d255-4492-bce4-bbebde6f2306",
   "metadata": {},
   "source": [
    "## Chunk References: Smaller Child Chunks Referring to Bigger Parent Chunk\n",
    "\n",
    "In this usage example, we show how to build a graph of smaller chunks pointing to bigger parent chunks.\n",
    "\n",
    "During query-time, we retrieve smaller chunks, but we follow references to bigger chunks. This allows us to have more context for synthesis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "49c784d8-71e6-42bc-84d9-a2aea4217b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_chunk_sizes = [128, 256, 512]\n",
    "sub_node_parsers = [\n",
    "    SimpleNodeParser.from_defaults(chunk_size=c) for c in sub_chunk_sizes\n",
    "]\n",
    "\n",
    "all_nodes = []\n",
    "for base_node in base_nodes:\n",
    "    for n in sub_node_parsers:\n",
    "        sub_nodes = n.get_nodes_from_documents([base_node])\n",
    "        sub_inodes = [\n",
    "            IndexNode.from_text_node(sn, base_node.node_id) for sn in sub_nodes\n",
    "        ]\n",
    "        all_nodes.extend(sub_inodes)\n",
    "\n",
    "    # also add original node to node\n",
    "    base_inode = IndexNode.from_text_node(base_node, base_node.node_id)\n",
    "    all_nodes.append(base_inode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2d614088-b122-40ad-811a-29cc0c2a295e",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_nodes_dict = {n.node_id: n for n in all_nodes}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a44ef2d5-0342-4073-831f-f35dd6f04dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_index_chunk = VectorStoreIndex(all_nodes, service_context=service_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c06af99f-02be-4055-a6ea-3071ffe8fc8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4c7c5e43-45b5-42d6-afc5-cb81ed3cb211",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever_chunk = RecursiveRetriever(\n",
    "    \"vector\",\n",
    "    retriever_dict={\"vector\": vector_retriever_chunk},\n",
    "    node_dict=all_nodes_dict,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9e9f7bcb-5442-4d2d-a7eb-814b68ebb45c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mRetrieving with query id None: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-21\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-21: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-22\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-22: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "**Node ID:** node-21<br>**Similarity:** 0.8693572142828401<br>**Text:** 22\n",
       "\n",
       "TruthfulQA ↑ToxiGen ↓\n",
       "MPT7B 29.13 22.32\n",
       "30B 35.25 22.61\n",
       "Falcon7B 25.95 14.53\n",
       "40B 40.39 23.44\n",
       "Llama 17B 27.42 23.00\n",
       "13B 41.74 23.08\n",
       "33B 44.19 22.57\n",
       "65B 48.71 21.77\n",
       "Llama 27B 33.29 21.25\n",
       "13B 41.86 26.10\n",
       "34B 43.45 21.19\n",
       "70B 50.18 24.60\n",
       "Table 11: Evaluation of pretrained LLMs on automatic safety benchmarks.For TruthfulQA, we present the\n",
       "percentageofgenerationsthatarebothtruthfulandinformative(thehigherthebetter).ForToxiGen,we\n",
       "present the percentage of toxic generations (the smaller, the better).Benchmarks give a summary view ofmodel capabilities and behaviors that allow us to understand general\n",
       "patternsinthemodel,buttheydonotprovideafullycomprehensiveviewoftheimpactthemodelmayhave\n",
       "onpeopleorreal-worldoutcomes;thatwouldrequirestudyofend-to-endproductdeployments.Further\n",
       "testing and mitigation should be done to understand bias and other social issues for the specific context\n",
       "in which a system may be deployed.For this, it may be necessary to test beyond the groups available in\n",
       "theBOLDdataset(race,religion,andgender).AsLLMsareintegratedanddeployed,welookforwardto\n",
       "continuing research that will amplify their potential for positive impact on these important social issues.4.2 Safety Fine-Tuning\n",
       "In this section, we describe our approach to safety fine-tuning, including safety categories, annotation\n",
       "guidelines,andthetechniquesweusetomitigatesafetyrisks.Weemployaprocesssimilartothegeneral\n",
       "fine-tuning methods as described in Section 3, with some notable differences related to safety concerns.Specifically, we use the following techniques in safety fine-tuning:\n",
       "1.Supervised Safety Fine-Tuning : We initialize by gathering adversarial prompts and safe demonstra-\n",
       "tions that are then included in the general supervised fine-tuning process (Section 3.1).This teaches\n",
       "themodeltoalignwithoursafetyguidelinesevenbeforeRLHF,andthuslaysthefoundationfor\n",
       "high-quality human preference data annotation.2.Safety RLHF : Subsequently, we integrate safety in the general RLHF pipeline described in Se...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Node ID:** node-22<br>**Similarity:** 0.854811252011206<br>**Text:** Wethendefinebestpracticesforsafeandhelpfulmodelresponses: themodelshouldfirstaddressimmediate\n",
       "safetyconcernsifapplicable,thenaddressthepromptbyexplainingthepotentialriskstotheuser,andfinally\n",
       "provide additional information if possible.We also ask the annotators to avoid negative user experience\n",
       "categories (see Appendix A.5.2).The guidelines are meant to be a general guide for the model and are\n",
       "iteratively refined and revised to include newly identified risks.4.2.2 Safety Supervised Fine-Tuning\n",
       "InaccordancewiththeestablishedguidelinesfromSection4.2.1,wegatherpromptsanddemonstrations\n",
       "ofsafemodelresponsesfromtrainedannotators,andusethedataforsupervisedfine-tuninginthesame\n",
       "manner as described in Section 3.1.An example can be found in Table 5.The annotators are instructed to initially come up with prompts that they think could potentially induce\n",
       "themodel toexhibit unsafebehavior, i.e.,perform redteaming, asdefined bythe guidelines.Subsequently,\n",
       "annotators are tasked with crafting a safe and helpful response that the model should produce.4.2.3 Safety RLHF\n",
       "Weobserveearlyinthedevelopmentof Llama 2-Chat thatitisabletogeneralizefromthesafedemonstrations\n",
       "insupervisedfine-tuning.Themodelquicklylearnstowritedetailedsaferesponses,addresssafetyconcerns,\n",
       "explainwhythetopicmightbesensitive,andprovideadditionalhelpfulinformation.Inparticular,when\n",
       "the model outputs safe responses, they are often more detailed than what the average annotator writes.Therefore, after gathering only a few thousand supervised demonstrations, we switched entirely to RLHF to\n",
       "teachthemodelhowtowritemorenuancedresponses.ComprehensivetuningwithRLHFhastheadded\n",
       "benefit that it may make the model more robust to jailbreak attempts (Bai et al., 2022a).WeconductRLHFbyfirstcollectinghumanpreferencedataforsafetysimilartoSection3.2.2: annotators\n",
       "writeapromptthattheybelievecanelicitunsafebehavior,andthencomparemultiplemodelresponsesto\n",
       "theprompts,selectingtheresponsethatissafestaccordingtoasetofguidelines.Wethenusethehu...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nodes = retriever_chunk.retrieve(\n",
    "    \"Can you tell me about the key concepts for safety finetuning\"\n",
    ")\n",
    "for node in nodes:\n",
    "    display_source_node(node, source_length=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "411f26ad-d13b-4858-938e-efcfa899e8cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine_chunk = RetrieverQueryEngine.from_args(\n",
    "    retriever_chunk, service_context=service_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4cd98366-0d5f-4d04-87cd-b811990b7485",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mRetrieving with query id None: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-21\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-21: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-22\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-22: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0mThe key concepts for safety fine-tuning include supervised safety fine-tuning, safety RLHF (Reinforcement Learning from Human Feedback), and safety context distillation. In supervised safety fine-tuning, adversarial prompts and safe demonstrations are gathered and included in the fine-tuning process to align the model with safety guidelines. Safety RLHF involves training a safety-specific reward model and gathering challenging adversarial prompts to optimize the model's safety through rejection sampling and PPO (Proximal Policy Optimization) optimization. Safety context distillation refines the RLHF pipeline by generating safer model responses using a safety preprompt and fine-tuning the model on these safer responses. These techniques aim to mitigate safety risks and improve the model's alignment with safety guidelines.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine_chunk.query(\n",
    "    \"Can you tell me about the key concepts for safety finetuning\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bcc7379-c077-40b7-ba4e-f47f80def0c7",
   "metadata": {},
   "source": [
    "## Metadata References: Summaries + Generated Questions referring to a bigger chunk\n",
    "\n",
    "In this usage example, we show how to define additional context that references the source node.\n",
    "\n",
    "This additional context includes summaries as well as generated questions.\n",
    "\n",
    "During query-time, we retrieve smaller chunks, but we follow references to bigger chunks. This allows us to have more context for synthesis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "24e40c5e-4868-487f-aaf4-f333aa4bda66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.node_parser import SimpleNodeParser\n",
    "from llama_index.schema import IndexNode\n",
    "from llama_index.node_parser.extractors import (\n",
    "    SummaryExtractor,\n",
    "    QuestionsAnsweredExtractor,\n",
    "    MetadataExtractor,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5c5d6f87-790e-4b82-abb2-cc6944678b00",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_extractor = MetadataExtractor(\n",
    "    extractors=[\n",
    "        SummaryExtractor(summaries=[\"self\"], show_progress=True),\n",
    "        QuestionsAnsweredExtractor(questions=5, show_progress=True),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e47c706c-940e-499d-b742-eaf09a230b0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run metadata extractor across base nodes, get back dictionaries\n",
    "metadata_dicts = metadata_extractor.extract(base_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2873327d-420a-4778-a83b-6fdf7aa21bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cache metadata dicts\n",
    "def save_metadata_dicts(path):\n",
    "    with open(path, \"w\") as fp:\n",
    "        for m in metadata_dicts:\n",
    "            fp.write(json.dumps(m) + \"\\n\")\n",
    "\n",
    "\n",
    "def load_metadata_dicts(path):\n",
    "    with open(path, \"r\") as fp:\n",
    "        metadata_dicts = [json.loads(l) for l in fp.readlines()]\n",
    "        return metadata_dicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "e318efb2-9afa-4414-b37f-71738d73d01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_metadata_dicts(\"data/llama2_metadata_dicts.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4edce99f-8a96-4539-95e7-62aeeabb2ce9",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_dicts = load_metadata_dicts(\"data/llama2_metadata_dicts.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f18d2109-5fcb-4fd5-b147-23897fed8787",
   "metadata": {},
   "outputs": [],
   "source": [
    "# all nodes consists of source nodes, along with metadata\n",
    "all_nodes = base_nodes\n",
    "for idx, d in enumerate(metadata_dicts):\n",
    "    inode_q = IndexNode(\n",
    "        text=d[\"questions_this_excerpt_can_answer\"], index_id=base_nodes[idx].node_id\n",
    "    )\n",
    "    inode_s = IndexNode(text=d[\"section_summary\"], index_id=base_nodes[idx].node_id)\n",
    "    all_nodes.extend([inode_q, inode_s])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "8f90ada6-0969-40cc-a4ec-3579b4900cdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_nodes_dict = {n.node_id: n for n in all_nodes}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "22abc768-83d5-41d0-84f0-533899c76894",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load index into vector index\n",
    "from llama_index import VectorStoreIndex, ServiceContext\n",
    "from llama_index.llms import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm)\n",
    "\n",
    "vector_index_metadata = VectorStoreIndex(all_nodes, service_context=service_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "0d53938a-1322-41b1-ad11-169b13b9805a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_retriever_metadata = vector_index_metadata.as_retriever(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "37ae791f-c183-4ad4-9a3a-253288ded5a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever_metadata = RecursiveRetriever(\n",
    "    \"vector\",\n",
    "    retriever_dict={\"vector\": vector_retriever_metadata},\n",
    "    node_dict=all_nodes_dict,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "3cd85685-19eb-44cc-ad27-1d163eaddad6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mRetrieving with query id None: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-21\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-21: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-27\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-27: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "**Node ID:** node-21<br>**Similarity:** 0.8455204474983238<br>**Text:** 22\n",
       "\n",
       "TruthfulQA ↑ToxiGen ↓\n",
       "MPT7B 29.13 22.32\n",
       "30B 35.25 22.61\n",
       "Falcon7B 25.95 14.53\n",
       "40B 40.39 23.44\n",
       "Llama 17B 27.42 23.00\n",
       "13B 41.74 23.08\n",
       "33B 44.19 22.57\n",
       "65B 48.71 21.77\n",
       "Llama 27B 33.29 21.25\n",
       "13B 41.86 26.10\n",
       "34B 43.45 21.19\n",
       "70B 50.18 24.60\n",
       "Table 11: Evaluation of pretrained LLMs on automatic safety benchmarks.For TruthfulQA, we present the\n",
       "percentageofgenerationsthatarebothtruthfulandinformative(thehigherthebetter).ForToxiGen,we\n",
       "present the percentage of toxic generations (the smaller, the better).Benchmarks give a summary view ofmodel capabilities and behaviors that allow us to understand general\n",
       "patternsinthemodel,buttheydonotprovideafullycomprehensiveviewoftheimpactthemodelmayhave\n",
       "onpeopleorreal-worldoutcomes;thatwouldrequirestudyofend-to-endproductdeployments.Further\n",
       "testing and mitigation should be done to understand bias and other social issues for the specific context\n",
       "in which a system may be deployed.For this, it may be necessary to test beyond the groups available in\n",
       "theBOLDdataset(race,religion,andgender).AsLLMsareintegratedanddeployed,welookforwardto\n",
       "continuing research that will amplify their potential for positive impact on these important social issues.4.2 Safety Fine-Tuning\n",
       "In this section, we describe our approach to safety fine-tuning, including safety categories, annotation\n",
       "guidelines,andthetechniquesweusetomitigatesafetyrisks.Weemployaprocesssimilartothegeneral\n",
       "fine-tuning methods as described in Section 3, with some notable differences related to safety concerns.Specifically, we use the following techniques in safety fine-tuning:\n",
       "1.Supervised Safety Fine-Tuning : We initialize by gathering adversarial prompts and safe demonstra-\n",
       "tions that are then included in the general supervised fine-tuning process (Section 3.1).This teaches\n",
       "themodeltoalignwithoursafetyguidelinesevenbeforeRLHF,andthuslaysthefoundationfor\n",
       "high-quality human preference data annotation.2.Safety RLHF : Subsequently, we integrate safety in the general RLHF pipeline described in Se...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Node ID:** node-27<br>**Similarity:** 0.8430926818606254<br>**Text:** In each case, we took the overall lessons as a guide\n",
       "to helpfurther modelsafetytraining, and specificallytook data fromthese exercisesformodel fine-tuning,\n",
       "model feedback training, and as a signal for other safety model training.Multiple additionalrounds ofred teaming wereperformed over severalmonths tomeasure the robustness\n",
       "of each new model as it was released internally.We defined the robustness of a model, γ, with respect to\n",
       "a red teaming exercise executed by a set of experts as the average number of created prompts that would\n",
       "triggeraviolatingresponsefromthemodelperpersonperhour.Asanexample,onour7Bmodel,wehadan\n",
       "evolution of γ: 1.8→0.45over several red teaming iterations and model refinements.Robustness will likely\n",
       "continue to improve with additional red teaming efforts.Another magnitude that we tracked as new models\n",
       "were produced was the percentage of prompts triggering violating responses discovered in the previous red\n",
       "teaming exercises that were mitigated in a given new candidate release.On average, we had a 90% rejection\n",
       "rate model over model.4.4 Safety Evaluation of Llama 2-Chat\n",
       "SafetyHumanEvaluation.Wecollectedroughly2,000adversarialpromptsforhumanevaluationaccording\n",
       "to the safety categories in Section 4.2.1, in which 1,351 prompts are single-turn and 623 are multi-turn.Examples of evaluation prompts and responses can be found in Appendix A.4.6.We then asked raters to judge models for safety violations on a five-point Likert scale with the following\n",
       "definitions:\n",
       "•5- No safety violations and very helpful\n",
       "•4- No safety violations and only minor non-safety issues\n",
       "•3- No safety violations but not helpful or other major non-safety issues\n",
       "•2- Mild or moderate safety violations\n",
       "29\n",
       "\n",
       "(a)Overall violation percentage.(b)Overall safety and helpfulness mean rating.Figure17: Overallsafetymeasures.Left: Llama 2-Chat haslowviolationpercentageoverallacrossmodel\n",
       "sizes.Right: Llama 2-Chat has high safety and helpfulness mean rating overall across model sizes.It is\n",
       "importan...<br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nodes = retriever_metadata.retrieve(\n",
    "    \"Can you tell me about the key concepts for safety finetuning\"\n",
    ")\n",
    "for node in nodes:\n",
    "    display_source_node(node, source_length=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "5285854a-69a6-4bc4-a2a5-1004cc790a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine_metadata = RetrieverQueryEngine.from_args(\n",
    "    retriever_metadata, service_context=service_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "4e0ada5c-9a83-4517-bbb7-899d4415d68a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mRetrieving with query id None: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-21\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-21: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: node-27\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id node-27: Can you tell me about the key concepts for safety finetuning\n",
      "\u001b[0mThe key concepts for safety fine-tuning include supervised safety fine-tuning, safety RLHF (Reinforcement Learning from Human Feedback), safety context distillation, safety categories, and annotation guidelines. Supervised safety fine-tuning involves gathering adversarial prompts and safe demonstrations to align the model with safety guidelines. Safety RLHF integrates safety into the RLHF pipeline by training a safety-specific reward model and gathering challenging adversarial prompts. Safety context distillation refines the RLHF pipeline by generating safer model responses through prefixing prompts with a safety preprompt. The safety categories considered are illicit and criminal activities, hateful and harmful activities, and unqualified advice. The annotation guidelines help create adversarial prompts along the dimensions of risk categories and attack vectors.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine_metadata.query(\n",
    "    \"Can you tell me about the key concepts for safety finetuning\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9973bdca-d179-47d6-bd96-2631b36e1d94",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "We evalate how well our recursive retrieval + node reference methods work. We evaluate both chunk references as well as metadata references. We use embedding similarity lookup to retrieve the reference nodes.\n",
    "\n",
    "We compare both methods against a baseline retriever where we fetch the raw nodes directly.\n",
    "\n",
    "In terms of metrics, we evaluate using both hit-rate and MRR."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3a30b7-2eb2-4eae-b0b9-1d4ec26ac915",
   "metadata": {},
   "source": [
    "### Dataset Generation\n",
    "\n",
    "We first generate a dataset of questions from the set of text chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "3fe8ae8a-a2b2-4515-bcff-1145e14ede3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.finetuning import (\n",
    "    generate_qa_embedding_pairs,\n",
    "    EmbeddingQAFinetuneDataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "eef1b43d-996b-4b0a-becb-1cec08d9f8c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 231/231 [19:05<00:00,  4.96s/it]\n"
     ]
    }
   ],
   "source": [
    "eval_dataset = generate_qa_embedding_pairs(base_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "bd3e2507-9157-48a5-909b-18eeb9ec01d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset.save_json(\"data/llama2_eval_dataset.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "611f07af-2006-4158-8dc6-59d11a269c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional\n",
    "eval_dataset = EmbeddingQAFinetuneDataset.from_json(\"data/llama2_eval_dataset.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "31ac1eb4-6bec-49ab-be7e-695e435fdcfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: generalize into eval functions\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def evaluate(\n",
    "    dataset,\n",
    "    retriever,\n",
    "    verbose=False,\n",
    "):\n",
    "    corpus = dataset.corpus\n",
    "    queries = dataset.queries\n",
    "    relevant_docs = dataset.relevant_docs\n",
    "\n",
    "    eval_results = []\n",
    "    for query_id, query in tqdm(queries.items()):\n",
    "        retrieved_nodes = retriever.retrieve(query)\n",
    "        retrieved_ids = [node.node.node_id for node in retrieved_nodes]\n",
    "        expected_id = relevant_docs[query_id][0]\n",
    "\n",
    "        # for tmp in retrieved_nodes:\n",
    "        #     print(f\"NODE: {tmp.node.node_id}\")\n",
    "        #     print(tmp.node.get_content())\n",
    "\n",
    "        rank = None\n",
    "        for idx, id in enumerate(retrieved_ids):\n",
    "            if id == expected_id:\n",
    "                rank = idx + 1\n",
    "                break\n",
    "\n",
    "        is_hit = rank is not None  # assume 1 relevant doc\n",
    "        mrr = 0 if rank is None else 1 / rank\n",
    "\n",
    "        eval_result = {\n",
    "            \"is_hit\": is_hit,\n",
    "            \"mrr\": mrr,\n",
    "            \"retrieved\": retrieved_ids,\n",
    "            \"expected\": expected_id,\n",
    "            \"query\": query_id,\n",
    "        }\n",
    "        eval_results.append(eval_result)\n",
    "\n",
    "    return eval_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb4782a6-f3da-453f-93be-7683ed15b508",
   "metadata": {},
   "source": [
    "### Compare Results\n",
    "\n",
    "We run evaluations on each of the retrievers to measure hit rate and MRR.\n",
    "\n",
    "We find that retrievers with node references (either chunk or metadata) tend to perform better than retrieving the raw chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "87798866-11bc-4f7f-b8aa-0a023309492f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# set vector retriever similarity top k to higher\n",
    "top_k = 10\n",
    "\n",
    "\n",
    "def display_results(names, results_arr):\n",
    "    \"\"\"Display results from evaluate.\"\"\"\n",
    "\n",
    "    hit_rates = []\n",
    "    mrrs = []\n",
    "    for name, results in zip(names, results_arr):\n",
    "        results_df = pd.DataFrame(results)\n",
    "        hit_rate = results_df[\"is_hit\"].mean()\n",
    "        mrr = results_df[\"mrr\"].mean()\n",
    "        hit_rates.append(hit_rate)\n",
    "        mrrs.append(mrr)\n",
    "\n",
    "    final_df = pd.DataFrame({\"retrievers\": names, \"hit_rate\": hit_rates, \"mrr\": mrrs})\n",
    "    display(final_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6d142c6-0374-43ec-af31-e02d246bd815",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=top_k)\n",
    "retriever_chunk = RecursiveRetriever(\n",
    "    \"vector\",\n",
    "    retriever_dict={\"vector\": vector_retriever_chunk},\n",
    "    node_dict=all_nodes_dict,\n",
    "    verbose=True,\n",
    ")\n",
    "results_chunk = evaluate(eval_dataset, retriever_chunk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae448fe7-3a66-45a6-8e8e-6ed3950e61b8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "vector_retriever_metadata = vector_index_metadata.as_retriever(similarity_top_k=top_k)\n",
    "retriever_metadata = RecursiveRetriever(\n",
    "    \"vector\",\n",
    "    retriever_dict={\"vector\": vector_retriever_metadata},\n",
    "    node_dict=all_nodes_dict,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "results_metadata = evaluate(eval_dataset, retriever_metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "3d3fc029-7ccc-4ec4-b391-b7b86744b5d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████| 509/509 [01:26<00:00,  5.88it/s]\n"
     ]
    }
   ],
   "source": [
    "base_retriever = base_index.as_retriever(similarity_top_k=top_k)\n",
    "results_base = evaluate(eval_dataset, base_retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "4ef0cd73-b1ad-4ec6-931f-357d2ceebd65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>retrievers</th>\n",
       "      <th>hit_rate</th>\n",
       "      <th>mrr</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Base Retriever</td>\n",
       "      <td>0.269155</td>\n",
       "      <td>0.191413</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Retriever (Chunk References)</td>\n",
       "      <td>0.292731</td>\n",
       "      <td>0.254551</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Retriever (Metadata References)</td>\n",
       "      <td>0.286837</td>\n",
       "      <td>0.240858</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        retrievers  hit_rate       mrr\n",
       "0                   Base Retriever  0.269155  0.191413\n",
       "1     Retriever (Chunk References)  0.292731  0.254551\n",
       "2  Retriever (Metadata References)  0.286837  0.240858"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_results(\n",
    "    [\n",
    "        \"Base Retriever\",\n",
    "        \"Retriever (Chunk References)\",\n",
    "        \"Retriever (Metadata References)\",\n",
    "    ],\n",
    "    [results_base, results_chunk, results_metadata],\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
